// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <boost/range/irange.hpp>
#include <melon/container/f14_set.h>
#include <pollux/type/simple_function_api.h>
#include <pollux/expression/eval_ctx.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/lib/lambda_function_util.h>
#include <pollux/functions/lib/rows_translation_util.h>

namespace kumo::pollux::functions {
    namespace {
        template<bool isMaxBy>
        class ArrayMinMaxByFunction : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /*outputType*/,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                POLLUX_CHECK_EQ(args.size(), 2);

                exec::LocalDecodedVector arrayDecoder(context, *args[0], rows);
                auto &decodedArray = *arrayDecoder.get();

                // Flatten input array.
                auto flatArray = flattenArray(rows, args[0], decodedArray);

                std::vector<VectorPtr> lambdaArgs = {flatArray->elements()};
                auto newNumElements = flatArray->elements()->size();

                SelectivityVector validRowsInReusedResult =
                        toElementRows<ArrayVector>(newNumElements, rows, flatArray.get());

                VectorPtr newElements;
                auto elementToTopLevelRows = getElementToTopLevelRows(
                    newNumElements, rows, flatArray.get(), context.pool());

                // Loop over lambda functions and apply these to elements of the base array.
                // In most cases there will be only one function and the loop will run once.
                auto it = args[1]->as_unchecked<FunctionVector>()->iterator(&rows);
                while (auto entry = it.next()) {
                    auto elementRows = toElementRows<ArrayVector>(
                        newNumElements, *entry.rows, flatArray.get());
                    auto wrapCapture = toWrapCapture<ArrayVector>(
                        newNumElements, entry.callable, *entry.rows, flatArray);

                    entry.callable->apply(
                        elementRows,
                        &validRowsInReusedResult,
                        wrapCapture,
                        &context,
                        lambdaArgs,
                        elementToTopLevelRows,
                        &newElements);
                }

                // New elements after the lambda function is applied
                exec::LocalDecodedVector decodedElements(
                    context, *newElements, validRowsInReusedResult);

                BufferPtr resultIndices = allocate_indices(rows.size(), context.pool());
                auto *rawResultIndices = resultIndices->asMutable<vector_size_t>();

                BufferPtr resultNulls = allocate_nulls(rows.size(), context.pool());
                auto *mutable_nulls = resultNulls->asMutable<uint64_t>();

                // Get the max/min value out of the output of the lambda
                auto decodedIndices = decodedElements->indices();
                const auto *baseElementsVector = decodedElements->base();

                CompareFlags flags{.nullsFirst = false, .ascending = isMaxBy};
                auto compareLogic = [&](vector_size_t a, vector_size_t b) {
                    if (a == b) {
                        return false;
                    }
                    bool aNull = decodedElements->is_null_at(a);
                    bool bNull = decodedElements->is_null_at(b);
                    if (aNull) {
                        return false;
                    }
                    if (bNull) {
                        return true;
                    }
                    std::optional<int32_t> result = baseElementsVector->compare(
                        baseElementsVector, decodedIndices[a], decodedIndices[b], flags);
                    // If lambda returns the same value, then return the last element
                    return result.value() <= 0;
                };

                context.applyToSelectedNoThrow(rows, [&](vector_size_t row) {
                    const auto size = flatArray->sizeAt(row);
                    const auto offset = flatArray->offsetAt(row);
                    const auto range = boost::irange(offset, offset + size);

                    int minMaxIndex =
                            *std::max_element(range.begin(), range.end(), compareLogic);
                    if (decodedElements->is_null_at(minMaxIndex)) {
                        bits::set_null(mutable_nulls, row);
                    }
                    rawResultIndices[row] = minMaxIndex;
                });

                auto localResult = BaseVector::wrap_in_dictionary(
                    resultNulls, resultIndices, rows.end(), flatArray->elements());
                context.moveOrCopyResult(localResult, rows, result);
            }
        };

        class ArrayMaxByFunction : public ArrayMinMaxByFunction<true> {
        };

        class ArrayMinByFunction : public ArrayMinMaxByFunction<false> {
        };
    } // namespace

    std::vector<std::shared_ptr<exec::FunctionSignature> > signature() {
        return {
            exec::FunctionSignatureBuilder()
            .typeVariable("T")
            .typeVariable("U")
            .returnType("T")
            .argumentType("array(T)")
            .argumentType("function(T, U)")
            .build()
        };
    }

    // Register function.
    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_array_max_by,
        signature(),
        std::make_unique<ArrayMaxByFunction>());

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_array_min_by,
        signature(),
        std::make_unique<ArrayMinByFunction>());
} // namespace kumo::pollux::functions
